#include "tunnel.h"
#include "base64.h"
#include "utils.h"

#include "muduo/base/ThreadLocal.h"
#include "muduo/base/AsyncLogging.h"
#include <cstdio>
#include <cstring>
#include <iostream>
#include <string>
#include <vector>
#include <fstream>
#include <unordered_map>
#include <cctype>
#include <filesystem>
#include <unistd.h>

off_t kRollSize = 500*1000*1000;

using namespace muduo;
using namespace muduo::net;
namespace fs = std::filesystem;

namespace {
static const int INTERVAL = 60;
static const int WR_INTERVAL = 30;
}

std::vector<InetAddress> g_backends;
ThreadLocal<std::map<string, TunnelPtr> > t_tunnels;
std::unordered_map<std::string, TcpConnectionPtr> allow_conn;

std::vector<std::string> userlist;
MutexLock g_mutex;
size_t g_current = 0;

const char *authreq_res = "HTTP/1.1 407 Proxy Authentication Required\r\n"
                          "Proxy-Authenticate: Basic realm=\"Access to internal site\"\r\n\r\n";

const char *unauth_res = "HTTP/1.1 401 Unauthorized\r\n"
                         "WWW-Authenticate: Basic realm=\"Access to staging site\"\r\n\r\n";

static std::vector<std::string> browser_agent_type = {
    "Mozilla",
    "Chrome",
    "Opera",
    "OPR",
    "Safari"
};

muduo::AsyncLogging* g_asyncLog = nullptr;

void asyncOutput(const char* msg, int len)
{
  g_asyncLog->append(msg, len);
}

bool CheckIsValidUser(const std::string &user)
{
  for (auto i : userlist) {
    if (user == i) {
      return true;
    }
  }

  return false;
}

void ReadUserList(const std::string &workDir)
{
  fs::path workPath(workDir);
  workPath.append("userlist.conf");
  LOG_INFO << "updated user list";
  std::fstream file(workPath.u8string(), file.in);
  std::vector<std::string> tmp;

  if (!file.is_open()) {
    LOG_ERROR << "open user list file failure";
    // insert an account for testing.
    tmp.push_back(utils::testing_account);
    userlist = tmp;
    return;
  }

  while (!file.eof()) {
    std::string user;
    std::getline(file, user);
    if (!user.empty()) {
      tmp.push_back(user);
    }
  }
  userlist = tmp;

  // insert an account for testing.
  userlist.push_back(utils::testing_account);
}

bool CheckIsBrowser(const std::string &head)
{
  std::string UserAgent("User-Agent");
  std::string::size_type agentCur = head.find(UserAgent);
  if (agentCur != std::string::npos) {
    std::string::size_type LFCR = head.find("\r\n", agentCur);
    if (LFCR != std::string::npos) {
      std::string agentType = head.substr(agentCur + UserAgent.size(),
                                          LFCR - agentCur - UserAgent.size());
      for (auto &i : browser_agent_type) {
        std::string::size_type tmpCur = agentType.find(i);
        if (tmpCur != std::string::npos) {
          return true;
        }
      }
    }
  }

  return false;
}

void onServerConnection(const TcpConnectionPtr& conn)
{
  LOG_INFO << (conn->connected() ? "UP" : "DOWN");
  std::map<string, TunnelPtr>& tunnels = t_tunnels.value();
  if (conn->connected())
  {
    conn->setTcpNoDelay(true);
    conn->stopRead();
    size_t current = 0;
    {
      MutexLockGuard guard(g_mutex);
      current = g_current;
      g_current = (g_current+1) % g_backends.size();
    }

    InetAddress backend = g_backends[current];
    TunnelPtr tunnel(new Tunnel(conn->getLoop(), backend, conn));
    tunnel->setup();
    tunnel->connect();
    tunnels[conn->name()] = tunnel;
  }
  else
  {
    assert(tunnels.find(conn->name()) != tunnels.end());
    tunnels[conn->name()]->disconnect();
    tunnels.erase(conn->name());
    allow_conn.erase(conn->name());
  }
}

void onServerMessage(const TcpConnectionPtr& conn, Buffer* buf, Timestamp)
{
  // just send raw-data that tls hand shaked conntions.
  std::string conn_name = conn->name();
  auto iter = allow_conn.find(conn_name);
  if (iter != allow_conn.end()) {
    iter->second->send(buf);
    return;
  }

  if (!conn->getContext().empty())
  {
    const TcpConnectionPtr& clientConn
      = boost::any_cast<const TcpConnectionPtr&>(conn->getContext());
    const char *begin = buf->peek();
    std::string buffer(begin);

    if (CheckIsBrowser(buffer)) {
      InetAddress peerAddr = conn->peerAddress();
      LOG_INFO << "There is a browser try to connect the program from " << peerAddr.toIpPort();
      conn->shutdown();
      return;
    }

    std::string::size_type auth_cur = buffer.find("Proxy-Authorization");
    if (auth_cur != std::string::npos) {
      std::string::size_type user_pass_cur = buffer.find("Basic", auth_cur);
      std::string::size_type endding = buffer.find("\r\n", user_pass_cur);
      std::string user_pass = buffer.substr(user_pass_cur + 6, endding - user_pass_cur - 6);
      std::string decode_user_pass = base64_decode(user_pass);
      if (CheckIsValidUser(decode_user_pass)) {
        clientConn->send(buf);
        allow_conn[conn->name()] = clientConn;
        std::map<string, TunnelPtr>& tunnels = t_tunnels.value();
        tunnels[conn->name()]->setUserName(decode_user_pass);
      } else {
        conn->send(unauth_res, strlen(unauth_res));
        conn->shutdown();
      }
    } else {
      conn->send(authreq_res, strlen(authreq_res));
      conn->shutdown();
    }
  }
}

int main(int argc, char* argv[])
{
  if (argc < 4)
  {
    fprintf(stderr, "Usage: %s confdir listen_port backend_ip:port [backend_ip:port]\n", argv[0]);
  }
  else
  {
    std::string workDir = argv[1];
    for (int i = 3; i < argc; ++i)
    {
      string hostport = argv[i];
      size_t colon = hostport.find(':');
      if (colon != string::npos)
      {
        string ip = hostport.substr(0, colon);
        uint16_t port = static_cast<uint16_t>(atoi(hostport.c_str()+colon+1));
        g_backends.push_back(InetAddress(ip, port));
      }
      else
      {
        fprintf(stderr, "invalid backend address %s\n", argv[i]);
        return 1;
      }
    }

    if (-1 == daemon(1, 0))
    {
      fprintf(stderr, "daemon error\n");
      exit(-1);
    }

    uint16_t port = static_cast<uint16_t>(atoi(argv[2]));
    InetAddress listenAddr(port);

    char name[256] = { '\0' };
    strncpy(name, argv[0], sizeof name - 1);
    muduo::AsyncLogging log(::basename(name), kRollSize);
    log.start();
    g_asyncLog = &log;
    muduo::Logger::setOutput(asyncOutput);
    ReadUserList(workDir);

    EventLoop loop;
    TcpServer server(&loop, listenAddr, "TcpBalancer");
    server.setConnectionCallback(onServerConnection);
    server.setMessageCallback(onServerMessage);
    server.setThreadNum(4);
    server.start();
    loop.runEvery(INTERVAL, std::bind(ReadUserList, workDir));
    loop.runEvery(WR_INTERVAL, std::bind(&Statistic::SyncStatisticData, Statistic::getInstance(), workDir));
    loop.loop();
  }
}


